package com.wy.test.protocol.saml2.saml20.provider.xml;

import java.util.ArrayList;
import java.util.HashMap;

import org.joda.time.DateTime;
import org.opensaml.Configuration;
import org.opensaml.saml2.core.Assertion;
import org.opensaml.saml2.core.AttributeStatement;
import org.opensaml.saml2.core.AuthnStatement;
import org.opensaml.saml2.core.Conditions;
import org.opensaml.saml2.core.Issuer;
import org.opensaml.saml2.core.Subject;
import org.opensaml.saml2.core.impl.AssertionBuilder;
import org.opensaml.xml.security.BasicSecurityConfiguration;
import org.opensaml.xml.security.credential.BasicCredential;
import org.opensaml.xml.security.keyinfo.KeyInfoGeneratorFactory;
import org.opensaml.xml.signature.Signature;
import org.opensaml.xml.signature.SignatureConstants;
import org.opensaml.xml.signature.Signer;
import org.opensaml.xml.signature.impl.SignatureBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;

import com.wy.test.authentication.core.web.AuthorizationUtils;
import com.wy.test.core.entity.UserEntity;
import com.wy.test.core.vo.AppSamlDetailVO;
import com.wy.test.core.web.AuthWebContext;
import com.wy.test.protocol.saml2.saml.service.IDService;
import com.wy.test.protocol.saml2.saml.service.TimeService;
import com.wy.test.protocol.saml2.saml20.binding.BindingAdapter;
import com.wy.test.protocol.saml2.saml20.xml.IssuerGenerator;

public class AssertionGenerator {

	private final static Logger logger = LoggerFactory.getLogger(AssertionGenerator.class);

	private final IssuerGenerator issuerGenerator;

	private final SubjectGenerator subjectGenerator;

	private final IDService idService;

	private final TimeService timeService;

	private final AuthnStatementGenerator authnStatementGenerator;

	private final AttributeStatementGenerator attributeStatementGenerator;

	private final ConditionsGenerator conditionsGenerator;

	public AssertionGenerator(String issuerName, TimeService timeService, IDService idService) {
		this.timeService = timeService;
		this.idService = idService;
		issuerGenerator = new IssuerGenerator(issuerName);
		subjectGenerator = new SubjectGenerator(timeService);
		authnStatementGenerator = new AuthnStatementGenerator();
		attributeStatementGenerator = new AttributeStatementGenerator();
		conditionsGenerator = new ConditionsGenerator();
	}

	public Assertion generateAssertion(AppSamlDetailVO saml20Details, BindingAdapter bindingAdapter,
			String assertionConsumerURL, String inResponseTo, String audienceUrl, int validInSeconds,
			HashMap<String, String> attributeMap, UserEntity userInfo) {

		Assertion assertion = new AssertionBuilder().buildObject();
		// Subject
		Subject subject = subjectGenerator.generateSubject(saml20Details, assertionConsumerURL, inResponseTo,
				validInSeconds, userInfo);
		assertion.setSubject(subject);
		// issuer
		Issuer issuer = issuerGenerator.generateIssuer();
		assertion.setIssuer(issuer);
		// AuthnStatements
		DateTime authnInstant = new DateTime(AuthWebContext.getSession().getCreationTime());
		AuthnStatement authnStatement = authnStatementGenerator.generateAuthnStatement(authnInstant);
		assertion.getAuthnStatements().add(authnStatement);
		// AttributeStatements
		ArrayList<GrantedAuthority> grantedAuthoritys = new ArrayList<GrantedAuthority>();
		grantedAuthoritys.add(new SimpleGrantedAuthority("ROLE_USER"));
		for (GrantedAuthority anthGrantedAuthority : ((UsernamePasswordAuthenticationToken) AuthorizationUtils
				.getAuthentication()).getAuthorities()) {
			grantedAuthoritys.add(anthGrantedAuthority);
		}
		AttributeStatement attributeStatement = attributeStatementGenerator.generateAttributeStatement(saml20Details,
				grantedAuthoritys, attributeMap, userInfo);
		assertion.getAttributeStatements().add(attributeStatement);
		// ID
		assertion.setID(idService.generateID());
		// IssueInstant
		assertion.setIssueInstant(timeService.getCurrentDateTime());
		// Conditions
		Conditions conditions = conditionsGenerator.generateConditions(audienceUrl, validInSeconds);
		assertion.setConditions(conditions);
		// sign Assertion
		try {
			if (bindingAdapter.getSigningCredential() == null) {
				throw new Exception("Signing Credential is null...");
			}
			logger.debug("EntityId " + bindingAdapter.getSigningCredential().getEntityId());
			BasicCredential basicCredential = new BasicCredential();
			basicCredential.setPrivateKey(bindingAdapter.getSigningCredential().getPrivateKey());

			Signature signature = new SignatureBuilder().buildObject();
			signature.setCanonicalizationAlgorithm(SignatureConstants.ALGO_ID_C14N_EXCL_OMIT_COMMENTS);

			signature.setSigningCredential(basicCredential);
			KeyInfoGeneratorFactory keyInfoGeneratorFactory = Configuration.getGlobalSecurityConfiguration()
					.getKeyInfoGeneratorManager().getDefaultManager().getFactory(bindingAdapter.getSigningCredential());

			signature.setKeyInfo(keyInfoGeneratorFactory.newInstance().generate(bindingAdapter.getSigningCredential()));
			BasicSecurityConfiguration config =
					(BasicSecurityConfiguration) Configuration.getGlobalSecurityConfiguration();

			if (saml20Details.getSignature().equalsIgnoreCase("RSAwithSHA1")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA1);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA1);
			} else if (saml20Details.getSignature().equalsIgnoreCase("RSAwithSHA256")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA256);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA256);
			} else if (saml20Details.getSignature().equalsIgnoreCase("RSAwithSHA384")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA384);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA384);
			} else if (saml20Details.getSignature().equalsIgnoreCase("RSAwithSHA512")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA512);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA512);
			} else if (saml20Details.getSignature().equalsIgnoreCase("RSAwithMD5")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_NOT_RECOMMENDED_RSA_MD5);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_NOT_RECOMMENDED_RSA_MD5);
			} else if (saml20Details.getSignature().equalsIgnoreCase("RSAwithRIPEMD160")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA_RIPEMD160);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_RIPEMD160);
			} else if (saml20Details.getSignature().equalsIgnoreCase("DSAwithSHA1")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA1);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA1);
			} else if (saml20Details.getSignature().equalsIgnoreCase("ECDSAwithSHA256")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA256);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA256);
			} else if (saml20Details.getSignature().equalsIgnoreCase("ECDSAwithSHA384")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA384);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA384);
			} else if (saml20Details.getSignature().equalsIgnoreCase("ECDSAwithSHA512")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA512);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_ECDSA_SHA512);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-MD5")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_NOT_RECOMMENDED_MD5);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_MAC_HMAC_NOT_RECOMMENDED_MD5);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-SHA1")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_SHA1);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_SIGNATURE_RSA_SHA1);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-SHA256")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_SHA256);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_MAC_HMAC_SHA256);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-SHA384")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_SHA384);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_MAC_HMAC_SHA384);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-SHA512")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_SHA512);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_MAC_HMAC_SHA512);
			} else if (saml20Details.getSignature().equalsIgnoreCase("HMAC-RIPEMD160")) {
				signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_MAC_HMAC_RIPEMD160);
				config.registerSignatureAlgorithmURI(saml20Details.getSignature(),
						SignatureConstants.ALGO_ID_MAC_HMAC_RIPEMD160);
			}

			if (saml20Details.getDigestMethod().equalsIgnoreCase("MD5")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_NOT_RECOMMENDED_MD5);
			} else if (saml20Details.getDigestMethod().equalsIgnoreCase("SHA1")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_SHA1);
			} else if (saml20Details.getDigestMethod().equalsIgnoreCase("SHA256")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_SHA256);
			} else if (saml20Details.getDigestMethod().equalsIgnoreCase("SHA384")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_SHA384);
			} else if (saml20Details.getDigestMethod().equalsIgnoreCase("SHA512")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_SHA512);
			} else if (saml20Details.getDigestMethod().equalsIgnoreCase("RIPEMD-160")) {
				config.setSignatureReferenceDigestMethod(SignatureConstants.ALGO_ID_DIGEST_RIPEMD160);
			}

			assertion.setSignature(signature);

			Configuration.getMarshallerFactory().getMarshaller(assertion).marshall(assertion);
			Signer.signObject(signature);

			logger.debug("assertion.isSigned " + assertion.isSigned());
		} catch (Exception e) {
			e.printStackTrace();
			logger.info("Unable to Signer assertion ", e);
		}

		return assertion;
	}
}
